import hashlib
from dataclasses import dataclass
from typing import List, Tuple, Optional, Union
from file_chunk_reader import FileChunkReader

@dataclass
class ChunkInfo:
    """数据块信息类"""
    chunk_number: int
    file_reads: List[Tuple[str, int, int]]  # (文件路径, 读取字节数, 起始位置)
    hash_raw_data: bytes  # 存储用于计算哈希值的原始数据
    digest: bytes  # 存储原始字节格式的哈希值

def print_chunk_info(chunk_info: ChunkInfo) -> None:
    """打印数据块信息
    
    Args:
        chunk_info (ChunkInfo): 数据块信息对象
    """
    print(f"\n数据块 #{chunk_info.chunk_number}")
    print("数据来源:")
    for file_path, bytes_read, start_pos in chunk_info.file_reads:
        print(f"  - 文件: {file_path}")
        print(f"    读取字节数: {bytes_read}")
        print(f"    起始位置: {start_pos}")
    print(f"HASH原文长度: {len(chunk_info.hash_raw_data)} 字节")
    print(f"SHA256: {chunk_info.digest.hex()}")  # 将字节转换为十六进制字符串

class ChunkHashCalculator:
    """数据块哈希计算器类"""
    def __init__(self, directory: str, chunk_size: int = 1024, max_files: int = 0):
        self.reader = FileChunkReader(
            directory=directory,
            chunk_size=chunk_size,
            max_files=max_files
        )
        self.chunk_counter = 0
    
    def reset(self) -> None:
        """重置计算器状态
        
        此方法将：
        1. 重置数据块计数器为0
        2. 重新初始化文件读取器
        """
        self.chunk_counter = 0
        self.reader.reset()
    
    def _prepare_hash_data(self, data: bytes, file_reads: List[Tuple[str, int, int]]) -> bytes:
        """准备用于计算哈希值的数据
        
        Args:
            data: 实际数据
            file_reads: 文件读取信息
            
        Returns:
            bytes: 用于计算哈希值的组合数据
        """
        # 首先添加计数器
        combined_data = f"{self.chunk_counter}".encode('utf-8')
        
        # 添加所有文件的信息和数据
        for file_path, bytes_read, start_pos in file_reads:
            # 添加文件信息
            file_info = f"{file_path}:{bytes_read}:{start_pos}".encode('utf-8')
            combined_data += file_info
        
        # 最后添加实际的数据块
        combined_data += data
        return combined_data
    
    def _process_next_chunk(self, return_chunk_info: bool = False) -> Optional[Union[bytes, ChunkInfo]]:
        """处理下一个数据块
        
        Args:
            return_chunk_info: 返回ChunkInfo信息
            
        Returns:
            Optional[Union[bytes, ChunkInfo]]: 如果return_chunk_info为True，返回ChunkInfo对象；
                                             如果为False，返回哈希值（字节格式）；
                                             如果读取失败，返回None
        """
        # 读取下一个数据块
        result = self.reader.read_next_chunk()
        if not result:
            return None
            
        data, file_reads = result
        self.chunk_counter += 1
        
        # 构建用于计算SHA256的输入数据
        combined_data = self._prepare_hash_data(data, file_reads)
        
        # 计算SHA256哈希值
        sha256_hash = hashlib.sha256(combined_data).digest()
        
        # 根据参数决定返回格式
        if return_chunk_info:
            return ChunkInfo(
                chunk_number=self.chunk_counter,
                file_reads=file_reads,
                hash_raw_data=combined_data,
                digest=sha256_hash
            )
        else:
            return sha256_hash
    
    def calculate_next_hash(self) -> Optional[bytes]:
        """计算下一个数据块的哈希值，只返回哈希值
        
        Returns:
            Optional[bytes]: 如果成功读取数据块，返回哈希值（字节格式）；
                           如果读取失败，返回None
        """
        result = self._process_next_chunk(return_chunk_info=False)
        return result if isinstance(result, bytes) else None
    
    def calculate_next_hash_with_info(self) -> Optional[ChunkInfo]:
        """计算下一个数据块的哈希值，返回完整信息
        
        Returns:
            Optional[ChunkInfo]: 如果成功读取数据块，返回包含完整信息的ChunkInfo对象；
                               如果读取失败，返回None
        """
        result = self._process_next_chunk(return_chunk_info=True)
        return result if isinstance(result, ChunkInfo) else None

def main():
    # 创建哈希计算器实例
    calculator = ChunkHashCalculator(
        directory=".",  # 当前目录
        chunk_size=4096,  # 每次读取1024字节
        max_files=0  # 读取所有文件
    )
    
    # 测试读取10个数据块
    print("测试完整信息接口:")
    for _ in range(5):
        chunk_info = calculator.calculate_next_hash_with_info()
        if chunk_info:
            print_chunk_info(chunk_info)
        else:
            print("读取数据失败")
            break
    
    calculator.reset()
    
    print("\n测试仅哈希值接口:")
    for _ in range(5):
        hash = calculator.calculate_next_hash()
        if hash:
            print(hash.hex())
        else:
            print("读取数据失败")
            break

if __name__ == "__main__":
    main() 